import asyncio
from dataclasses import dataclass
from typing import Literal, Optional
from agents import Agent, Runner, OpenAIProvider, RunResult, RunConfig, set_tracing_export_api_key, set_tracing_disabled
from openai import AsyncOpenAI, APIConnectionError
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import os


# 定义评估反馈数据结构
@dataclass
class EvaluationFeedback:
    score: Literal["pass", "needs_improvement", "fail"]
    feedback: str


# 配置API信息
BASE_URL = "https://api.chatfire.cn/v1"
API_KEY = "sk-UcrTOYgPf03qqxdGspmJtaUmP7OTkkzuxsRHWPpMSxKXK8IN"
os.environ["OPENAI_API_KEY"] = API_KEY

# 配置Tracing客户端使用相同的API密钥
set_tracing_export_api_key(API_KEY)
set_tracing_disabled(True)

# 创建故事生成Agent
story_outline_generator = Agent(
    name="story_outline_generator",
    model="gpt-4o",
    instructions=(
        "You generate a very short story outline based on the user's input."
        "If there is any feedback provided, use it to improve the outline."
    )
)

# 创建评估Agent
evaluator = Agent(
    name="evaluator",
    model="gpt-4o",
    instructions=(
        "You evaluate a story outline and decide if it's good enough."
        "If it's not good enough, you provide feedback on what needs improvement."
        "Never give it a pass on the first try."
    ),
    output_type=EvaluationFeedback
)

# 配置OpenAI服务
try:
    # 增加超时时间到30秒
    provider = OpenAIProvider(
        openai_client=AsyncOpenAI(
            base_url=BASE_URL,
            api_key=API_KEY,
            timeout=30.0  # 增加超时时间
        ),
        use_responses=False
    )
    print("✅ API客户端配置成功")
except Exception as e:
    print(f"❌ 配置API客户端时出错: {e}")
    exit(1)


# 带重试机制的运行函数
@retry(
    stop=stop_after_attempt(3),  # 最多重试3次
    wait=wait_exponential(multiplier=1, min=4, max=10),  # 指数退避策略
    retry=retry_if_exception_type((APIConnectionError, TimeoutError)),
)
async def safe_runner_run(agent, messages, provider):
    print(f"⚡ 正在运行Agent: {agent.name}...")
    return await Runner.run(agent, messages, run_config=RunConfig(model_provider=provider))


async def main() -> None:
    msg: str = input("What kind of story would you like to hear? ")
    messages = [{"content": msg, "role": "user"}]
    latest_outline: Optional[str] = None

    try:
        max_evaluations = 1  # 最大评估次数
        evaluation_count = 0  # 当前评估计数
        passed = False  # 是否通过评估的标志
        while evaluation_count < max_evaluations:
            evaluation_count += 1
            # 生成故事大纲
            try:
                story_outline_result: RunResult = await safe_runner_run(
                    story_outline_generator,
                    messages,
                    provider
                )
                latest_outline = story_outline_result.final_output
                print("\n=== Generated Story Outline ===")
                print(latest_outline)

                # 添加生成结果到消息历史
                messages.append({"role": "assistant", "content": latest_outline})
            except APIConnectionError as e:
                print(f"\n⚠️ 连接API服务器失败: {e}")
                return
            except Exception as e:
                print(f"\n❌ 生成故事时出错: {e}")
                return

            # 评估故事大纲
            try:
                print(f"\n=== 评估轮次 {evaluation_count}/{max_evaluations} ===")
                evaluator_result: RunResult = await safe_runner_run(
                    evaluator,
                    messages,
                    provider
                )
                result: EvaluationFeedback = evaluator_result.final_output
                print(f"\n=== Evaluation: {result.score} ===")
                print(result.feedback)

                # 检查是否通过评估
                if result.score == "pass":
                    print("\n✅ Story outline approved!")
                    break

                # 添加反馈到消息历史
                print("\n🔄 Re-running with feedback...")
                messages.append({
                    "role": "user",
                    "content": f"Evaluation feedback: {result.feedback}"
                })
            except Exception as e:
                print(f"\n❌ 评估故事时出错: {e}")
                return

        print(f"\n🌟 Final story outline:\n{latest_outline}")

    except KeyboardInterrupt:
        print("\n🚫 程序被用户中断")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except Exception as e:
        print(f"🔥 未处理的异常: {e}")